﻿"""
Jieun Oh

quantize.py -- Jieun's implementation of quantize for Music 422 HW2

"""
import numpy as np


#------------------------------Part 1: UNIFORM MIDTREAD QUANTIZATION-----------------------------#

def QuantizeUniform(aNum, nBits=16):
        """uniform mid-tread quantization of floating point number aNum with user-specified number of bits nBits"""
        if aNum < 0.0:
           sCode = 1<<(nBits-1)
        else:
           sCode = 0
           
        if abs(aNum)>=1.0:
           absCode = (1<<(nBits-1))-1
        else:
           absCode = np.int(( ((1<<nBits) -1) * abs(aNum) +1)/2)
			             
        aCode = sCode|absCode
        return aCode

     

def DequantizeUniform(aCode, nBits=16):
        """dequantization of uniform mid-tread code aCode with user-specified number of bits nBits"""
        if aCode >> (nBits-1):            # negative
           sign = -1
           aCode &= ~(1<<(nBits-1));      # clear most significant bit
        else:
           sign = 1

        absNumber = 2 * np.float(aCode) / ((1<<nBits)-1)
        aNum = sign*absNumber
        return aNum
        
       

def vQuantizeUniform(aNumVec, nBits=16):
        """performs QuantizeUniform for an array of numbers. aNumVec is a numpy array."""
        negative = aNumVec<0.0              # true if aNum in aNumVec is negative
        sCodeVec = negative<<(nBits-1)

        overflow = (abs(aNumVec)>=1.0)      # true if abs(aNum) >= 1.0
        
        absCodeVec = np.cast['int']((((1<<nBits)-1) * abs(aNumVec) + 1)/2)   # assume for now there's no overflow
        absCodeVec[overflow] = (1<<(nBits-1))-1                              # fix values for overflow elements

        aCodeVec = sCodeVec|absCodeVec
        return aCodeVec



def vDequantizeUniform(aCodeVec, nBits=16):
        """performs DequantizeUniform for an array of codes"""
        negative = np.cast['bool'](aCodeVec>>(nBits-1))         # true if aCode in aCodeVec is negative
        aCodeVec[negative] &= ~(1<<(nBits-1))  # if negative, clear most significant bit

        aNumberVec = 2 * (np.cast['float'](aCodeVec)) / ((1<<nBits)-1)
        aNumberVec[negative] *= -1.0           # if negative, multiply by -1

        return aNumberVec


#------------------------------Part 2: FLOATING POINT NON-BLOCK QUANTIZATION-----------------------------#

def ScaleFactorNonBlock(aNum, nScaleBits=3, nMantBits=5):
        """returns scale factor bits for floating point quantization (non-block)"""
        R = (1<<nScaleBits)-1+nMantBits
        aCode = QuantizeUniform(aNum, R)
   
        aCode &= ~(1<<(R-1))                     # clear the most significant digit (sign)

        numLeadingZeros = R-1
        while (aCode != 0):
            aCode = aCode>>1
            numLeadingZeros -= 1         

        if numLeadingZeros < (1<<nScaleBits)-1:
            scale = numLeadingZeros
        else:        
            scale = (1<<nScaleBits)-1
        return scale 


def MantissaNonBlock(aNum, scale, nScaleBits=3, nMantBits=5):
        """returns mantissa bits for floating point quantization (non-block)"""
        R = (1<<nScaleBits)-1+nMantBits
        aCode = QuantizeUniform(aNum, R)
        aCode &= ~(1<<(R-1))                     # clear the most significant digit (sign)

        if aNum < 0.0:
           sCode = 1<<(nMantBits-1)
        else:
           sCode = 0

        twoToRSminusOne = (1<<nScaleBits)-1
        if scale==twoToRSminusOne:
           restOfM = aCode          # guarenteed to be the (nMantBits-1) least significant digits of aCode
        else:
           # scale represents # of leading zeros
           numDigitsAfterLeadingOne = R-scale-2        # -1 for sign, -1 for leading one
           mask = (1<<numDigitsAfterLeadingOne)-1      # has numDigitsAfterLeadingOne 1's
           aCodeAfterLeadingOne = aCode & mask
           
           # take the (nMantBits-1) most significant digits of aCodeAfterLeadingOne
           restOfM = aCodeAfterLeadingOne>>(numDigitsAfterLeadingOne-(nMantBits-1))     
 
        mantissa = sCode|restOfM
        return mantissa


def DequantizeNonBlock(scale, mantissa, nScaleBits=3, nMantBits=5):
        """returns dequantized number for floating point quantizatino (non-block)."""
        R = (1<<nScaleBits)-1+nMantBits

        if mantissa >> (nMantBits-1):            # negative
           sCode = 1<<(R-1)
        else:
           sCode = 0

        twoToRSminusOne = (1<<nScaleBits)-1
        mask = (1<<(nMantBits-1))-1              # has nMantBits-1 1's
        restOfM = mantissa & mask

        if scale == twoToRSminusOne:
           # mantissa bits (excluding the sign bit) gets copied to the (nMantBits-1) least significant digits of R
           absCode = restOfM
        else: 
           # place the leading one in (1<< (R-scale-2) )
           # followed by (nMantBits-1) least significant digits of R 
           # followed by 1 0 0 ... 
           leadingOne = 1<<(R-scale-2)
           numSpare = R-scale-nMantBits-1        # to be filled with rounded values
           restOfM = restOfM<<numSpare
           if (numSpare > 0):
              spare = 1<<(numSpare-1)
              restOfM = restOfM|spare
           absCode = leadingOne|restOfM
 
        aCode = sCode|absCode

        aNum = DequantizeUniform(aCode, R)
        return aNum


#------------------------------Part 3: FLOATING POINT BLOCK QUANTIZATION-----------------------------#


def ScaleFactor(aNum, nScaleBits=3, nMantBits=5):
        """returns scale factor bits for floating point quantization (block)"""
        return ScaleFactorNonBlock(aNum, nScaleBits, nMantBits)


def Mantissa(aNum, scale, nScaleBits=3, nMantBits=5):
        """returns mantissa bits for floating point quantization (block). Can't assume a leading one after scale # of zeros"""
        R = (1<<nScaleBits)-1+nMantBits
        aCode = QuantizeUniform(aNum, R)
        aCode &= ~(1<<(R-1))                     # clear the most significant digit (sign)

        if aNum < 0.0:
           sCode = 1<<(nMantBits-1)
        else:
           sCode = 0

        # restOfM :=  (nMantBits-1) most significant digis from the (R-1-scale)least significant digits of aCode
        #             That is, shift aCode (R-scale-nMantBits) to the right
        restOfM = aCode>>(R-scale-nMantBits)

        mantissa = sCode|restOfM
		
        return mantissa


def Dequantize(scale, mantissa, nScaleBits=3, nMantBits=5):
        """returns dequantized number for floating point quantizatino (block). Can't assume a leading one after scale # of zeros"""
        R = (1<<nScaleBits)-1+nMantBits

        if mantissa >> (nMantBits-1):            # negative
           sCode = 1<<(R-1)
        else:
           sCode = 0

        mask = (1<<(nMantBits-1))-1              # has (nMantBits-1) 1's
        restOfM = mantissa & mask

	# absCode := (scale # 0's)+(restOfM)+[ (R-scale-nMantBits) of spare ]
        numSpare = R-scale-nMantBits             # to be filled with rounded values
        restOfM = restOfM<<numSpare
        if (numSpare > 0):
              if (restOfM != 0):                 # added
                 spare = 1<<(numSpare-1)
                 restOfM = restOfM|spare
 
        aCode = sCode|restOfM
        aNum = DequantizeUniform(aCode, R)

        return aNum


def vMantissa(aNumVec, scale, nScaleBits=3, nMantBits=5):
        """Returns vector of block-floating point mantissas for vector of signed fraction numbers given scale factor and numbers of scale and mantissa bits. Can't assume a leading one after scale # of zeros"""
        R = (1<<nScaleBits)-1+nMantBits
        aCodeVec = vQuantizeUniform(aNumVec, R)
        aCodeVec &= ~(1<<(R-1))                       # clear the most significant digit (sign)

        negative = aNumVec<0.0                        # true if aNum in aNumVec is negative
        sCodeVec = negative<<(nMantBits-1)

        restOfMVec = aCodeVec>>(R-scale-nMantBits)

        mantissaVec = sCodeVec|restOfMVec

        return mantissaVec


def vDequantize(scale, mantissaVec, nScaleBits=3, nMantBits=5):
        """Returns vector of signed fraction numbers for block floating point scale and vector of mantissas mantissa given specified scale and mantissa bits."""
        R = (1<<nScaleBits)-1+nMantBits

        negative = np.cast['bool'](mantissaVec>>(nMantBits-1))         # true if negative
        sCodeVec = negative<<(R-1)

        mask = (1<<(nMantBits-1))-1                   # has (nMantBits-1) 1's
        restOfMVec = mantissaVec & mask
        
        numSpare = R-scale-nMantBits
        restOfMVec = restOfMVec<<numSpare

        restOfMisZero = np.cast['int'](restOfMVec!=0)                 # true if restOfM=0, added so that no "1" used for rounding if mantissa==0

        if (numSpare > 0):
              spare = (restOfMisZero)*(1<<(numSpare-1))
              restOfMVec = restOfMVec|spare

        aCodeVec = sCodeVec|restOfMVec
        aNumVec = vDequantizeUniform(aCodeVec, R)

        return aNumVec



def smallestScale(subBlock, nScaleBits=3, nMantBits=5):
        """helper function for carrying out block floating point quantization: determine the scale given a subblock"""
        '''
        smallestScale = (1<<nScaleBits)-1
        for elem in subBlock:
            scale = ScaleFactor(elem, nScaleBits, nMantBits)
            if scale<smallestScale:
                 smallestScale = scale
        return smallestScale
        '''
        biggestElem = 0
        for elem in subBlock:
            if abs(elem) > biggestElem:
               biggestElem = abs(elem)
        return ScaleFactor(biggestElem, nScaleBits, nMantBits)



